
from fastapi import FastAPI, Request
from fastapi.middleware.cors import CORSMiddleware
from fastapi.middleware.trustedhost import TrustedHostMiddleware
from fastapi.middleware.gzip import GZipMiddleware
from starlette.middleware.base import BaseHTTPMiddleware
from config import get_config


class CustomHeaderMiddleware(BaseHTTPMiddleware):
    async def dispatch(self, req: Request, call_next:callable):
        response = await call_next(req)
        return response
    
def register_middleware(app: FastAPI):
    """
    注册中间件
    """
    settings = get_config()
    # 跨域
    app.add_middleware(
        CORSMiddleware,
        allow_origins=settings.CORS_ORIGINS,
        allow_credentials=settings.CORS_ALLOW_CREDENTIALS,
        allow_methods=settings.CORS_ALLOW_METHODS,
        allow_headers=settings.CORS_ALLOW_HEADERS,
    )
    # 允许的 hosts 列表
    app.add_middleware(TrustedHostMiddleware, allowed_hosts=settings.AllowHosts)
    # 处理包含”gzip”在Accept-Encoding标头中的任何请求的 GZip 响应。
    app.add_middleware(GZipMiddleware, minimum_size=1000)

    # 自定义中间件
    app.add_middleware(CustomHeaderMiddleware)